import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image


def add_gaussian_noise(image, mean=0, sigma=25):
    """
    Add Gaussian noise to an image
    
    Args:
        image: Input image
        mean: Mean of the Gaussian noise
        sigma: Standard deviation of the Gaussian noise
        
    Returns:
        Image with Gaussian noise
    """
    row, col, ch = image.shape
    gauss = np.random.normal(mean, sigma, (row, col, ch))
    gauss = gauss.reshape(row, col, ch)
    noisy = image + gauss
    noisy = np.clip(noisy, 0, 255)
    return noisy.astype(np.uint8)


def add_salt_pepper_noise(image, salt_prob=0.02, pepper_prob=0.02):
    """
    Add salt and pepper noise to an image
    
    Args:
        image: Input image
        salt_prob: Probability of adding salt noise
        pepper_prob: Probability of adding pepper noise
        
    Returns:
        Image with salt and pepper noise
    """
    noisy = np.copy(image)
    
    # Salt noise (white)
    salt_mask = np.random.random(image.shape) < salt_prob
    noisy[salt_mask] = 255
    
    # Pepper noise (black)
    pepper_mask = np.random.random(image.shape) < pepper_prob
    noisy[pepper_mask] = 0
    
    return noisy


def add_combined_noise(image, mean=0, sigma=25, salt_prob=0.01, pepper_prob=0.01):
    """
    Add both Gaussian and salt-and-pepper noise to an image
    
    Args:
        image: Input image
        mean: Mean of the Gaussian noise
        sigma: Standard deviation of the Gaussian noise
        salt_prob: Probability of adding salt noise
        pepper_prob: Probability of adding pepper noise
        
    Returns:
        Image with combined noise
    """
    # First add Gaussian noise
    gaussian_noisy = add_gaussian_noise(image, mean, sigma)
    
    # Then add salt and pepper noise to the result
    combined_noisy = add_salt_pepper_noise(gaussian_noisy, salt_prob, pepper_prob)
    
    return combined_noisy


def display_images(images, titles):
    """
    Display multiple images with their titles
    
    Args:
        images: List of images to display
        titles: List of titles for each image
    """
    plt.figure(figsize=(15, 5))
    
    for i, (image, title) in enumerate(zip(images, titles)):
        plt.subplot(1, len(images), i + 1)
        
        # Convert BGR to RGB for display with matplotlib
        if len(image.shape) == 3 and image.shape[2] == 3:
            plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        else:
            plt.imshow(image, cmap='gray')
            
        plt.title(title)
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()


def main():
    # Load the image
    image_path = 'pic1.jpg'
    original_image = cv2.imread(image_path)
    
    if original_image is None:
        print(f"Error: Could not load image from {image_path}")
        return
    
    # Add Gaussian noise with reduced sigma (less noise)
    gaussian_noise_image = add_gaussian_noise(original_image, mean=0, sigma=10)
    
    # Add salt and pepper noise with reduced probabilities
    salt_pepper_noise_image = add_salt_pepper_noise(original_image, salt_prob=0.005, pepper_prob=0.005)
    
    # Add combined noise with reduced parameters
    combined_noise_image = add_combined_noise(original_image, mean=0, sigma=8, salt_prob=0.003, pepper_prob=0.003)
    
    # Display the results
    display_images(
        [original_image, gaussian_noise_image, salt_pepper_noise_image, combined_noise_image],
        ['原始图像', '高斯噪声', '椒盐噪声', '混合噪声']
    )
    
    # Save the processed images
    cv2.imwrite('gaussian_noise.jpg', gaussian_noise_image)
    cv2.imwrite('salt_pepper_noise.jpg', salt_pepper_noise_image)
    cv2.imwrite('combined_noise.jpg', combined_noise_image)
    print("处理后的图像已保存为 'gaussian_noise.jpg'、'salt_pepper_noise.jpg' 和 'combined_noise.jpg'")


if __name__ == "__main__":
    main() 